package com.kklsqm.langchain4j.controller;

import com.kklsqm.langchain4j.assistant.DeepSeekAssistant;
import com.kklsqm.langchain4j.dto.ChatRequest;
import com.kklsqm.langchain4j.dto.ChatResponse;
import com.kklsqm.langchain4j.service.impl.LangChainTTSServiceImpl;

import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;
import reactor.core.publisher.Flux;

/**
 * 功能:
 * 作者: 沙琪马
 * 日期: 2025/6/3 17:07
 */
@Controller
public class ChatWebSocketController {

    @Resource
    private LangChainTTSServiceImpl chatService;

    @Resource
    private SimpMessagingTemplate messagingTemplate;

    @Resource
    private DeepSeekAssistant deepSeekAssistant;


    /**
     * 处理WebSocket聊天请求
     * @param request 包含聊天参数的请求对象
     */
    @MessageMapping("/chat.stream")
    public void handleChatStream(ChatRequest request) {
        Flux<String> responseFlux = chatService.chatStreamFlux(
                deepSeekAssistant,
                request.getSessionId(),
                request.getRole(),
                request.getQuestion(),
                request.getExtraInfo(),
                "deepseek-chatt"
        );

        // 订阅Flux并将结果发送到WebSocket
        responseFlux.subscribe(chunk ->
                        messagingTemplate.convertAndSend(
                                "/topic/chat.stream/" + request.getSessionId(),
                                new ChatResponse(chunk, false)
                        ),
                error ->
                        messagingTemplate.convertAndSend(
                                "/topic/chat.stream/" + request.getSessionId(),
                                new ChatResponse("Error: " + error.getMessage(), true)
                        ),
                () ->
                        messagingTemplate.convertAndSend(
                                "/topic/chat.stream/" + request.getSessionId(),
                                new ChatResponse("", true)
                        )
        );
    }
}